import torch
import numpy as np
import torch.nn as nn
import torch.nn.functional as F


class TestNet(nn.Module):
    def __init__(self):
        super(TestNet, self).__init__()
        self.op = nn.GRU(10, 16, 2, batch_first=True, bidirectional=False)

    def forward(self, x, h0):
        out, hn = self.op(x, h0)
        return [out]


if __name__ == "__main__":
    net = TestNet()

    x = torch.FloatTensor(np.ones([3, 5, 10]).astype(np.float32))
    h0 = torch.FloatTensor(np.ones([2, 3, 16]).astype(np.float32))
    result = net(x, h0)
    print(result[0])
    print(result[0].mean())
    print(result[0].shape)


